/*
 * Project 1 Module
 * Leapfrog Pseudo-Random Number Generator
 * Read seed and write random numbers to /proc/lfprng
 * Such that all threads from the same process have the same original seed
 */
#include <linux/init.h>
#include <linux/module.h>
#include <linux/proc_fs.h>
#include <linux/kernel.h>
#include <linux/sched.h>
#include <linux/vmalloc.h>
#include <linux/string.h>
#include <linux/list.h>
#include <asm/uaccess.h>
#include <asm/current.h>
#include <linux/cdev.h>


/*
 * Proc Functions and Variables
 */
static struct proc_dir_entry *proc_entry;

static int proc_read_lfprng(char *page, char **start,
			    off_t off, int count,
			    int *eoof, void *data);

static int proc_write_lfprng(struct file *file,
			     const char *buf,
			     unsigned long count,
			     void *data);

/*
 * LFPRNG Functions and Variables
 */
static const long long MULTIPLIER = 764261123;
static const long long PMOD = 2147483647;

struct ProcSeeds
{
  int pid;
  long long seed;

  struct list_head list;
};

struct TGIDS
{
  long long init_seed;
  int tgid;
  struct list_head procs;

  struct list_head list;
};

LIST_HEAD(thread_groups);

static void writeSeed(int tgid, int pid, long long seed);
static void writeSeed_w_tg(struct TGIDS *thread_group, int pid, long long seed);

static long long readRnd(int tgid, int pid);
static long long readRnd_w_tg(struct TGIDS *thread_group, int pid);

static struct TGIDS *makeThreadGroup(int tgid);
static struct ProcSeeds *makeProc(int pid);

static struct TGIDS *getTGID(int tgid);
static struct ProcSeeds *getProc(struct TGIDS *thread_group, int pid);

static void addProc(struct TGIDS *thread_group, struct ProcSeeds *proc);

static void ReseedProcList(struct TGIDS *thread_group, long long seed);
static long long getRand(long long seed);

static long long getProcsRandom(struct ProcSeeds *proc);

static void dumpProcList(struct TGIDS *thread_group);

/*
 * Kernel Functions
 */
static int lfprng_init(void)
{
  proc_entry = create_proc_entry("lfprng",
				 0666, NULL);

  if(!proc_entry)
    {
      return -ENOMEM;

    }

  proc_entry->read_proc = proc_read_lfprng;
  proc_entry->write_proc = proc_write_lfprng;

  printk(KERN_INFO "STARTING NEW LFPRNG MODULE----------------------------\n");

  return 0;
}

static void lfprng_exit(void)
{
  remove_proc_entry("lfprng",NULL);
}


/*
 * Proc Functions
 */
static int proc_read_lfprng(char *buf, char **start,
			    off_t off, int count,
			    int *eof, void *data)
{
  long long val;
  int len;

  printk(KERN_INFO "GET RAND, TGID: %d, PID: %d\n", current->tgid, current->pid);
  val = readRnd(current->tgid, current->pid);

  len = sprintf(buf, "%lld", val);

  return len;
}

static int proc_write_lfprng(struct file *file,
			     const char *buf,
			     unsigned long count,
			     void *data)
{
  static const int MAX_SIZE = sizeof(long long);
  long long seed;

  printk(KERN_INFO "SEEDING, TGID: %d, PID: %d\n", current->tgid, current->pid);
  if(count > MAX_SIZE)
    {
      count = MAX_SIZE;
    }

  if(copy_from_user(&seed, buf, count))
    {
      return -EFAULT;
    }
  /*
  printk(KERN_INFO "seed contains: \"%lld\"\n", seed);
  */
  writeSeed(current->tgid, current->pid, seed);

  return count;
}

/*
 * LFPRNG Functions
 */
static void writeSeed(int tgid, int pid, long long seed)
{
  struct TGIDS *thread_group;

  thread_group = getTGID(tgid);

  /* Is this a new thread group? */
  if(thread_group == NULL)
    { /* yes, add it to the list */
      /*
      printk(KERN_INFO "New Thread Group\n");
      */
      thread_group = makeThreadGroup(tgid);
      list_add(&thread_group->list, &thread_groups);
    }
  /*
  printk(KERN_INFO "Writing Seed to Thread Group %d\n", thread_group->tgid);
  */
  writeSeed_w_tg(thread_group, pid, seed);
}

static void writeSeed_w_tg(struct TGIDS *thread_group, int pid, long long seed)
{
  struct ProcSeeds *proc;

  proc = getProc(thread_group, pid);

  /* Is this a new Proc? */
  if(proc == NULL)
    { /* yes, add it to the list */
      /*
 printk(KERN_INFO "New Proc\n");
      */
      proc = makeProc(pid);
      list_add(&proc->list, &thread_group->procs);
    }

  /*
  printk(KERN_INFO "Reseeding with seed: %lld for thread group %d at proc %d\n", seed, thread_group->tgid, proc->pid);
  */
  ReseedProcList(thread_group, seed);
}

static long long readRnd(int tgid, int pid)
{
  struct TGIDS *thread_group;

  thread_group = getTGID(tgid);

  /* Is this a new thread group? */
  if(thread_group == NULL)
    { /* Yes, seed it */
      /*
      printk(KERN_INFO "New Thread Group, seed it\n");
      */
      writeSeed(tgid, pid, ((MULTIPLIER*pid)%PMOD));

      /* lets try it again */
      thread_group = getTGID(tgid);
    }

  /*
  printk(KERN_INFO "Reading rand from thread group %d\n", thread_group->tgid);
  */
  return readRnd_w_tg(thread_group, pid);
}

static long long readRnd_w_tg(struct TGIDS *thread_group, int pid)
{
  struct ProcSeeds *proc;

  proc = getProc(thread_group, pid);

  if(proc == NULL)
    {
      /*
      printk(KERN_INFO "New Proc!\n");
      */
      proc = makeProc(pid);
      addProc(thread_group, proc);
    }
  /*
  printk(KERN_INFO "Reading rand from proc %d\n", proc->pid);
  */
  return getProcsRandom(proc);
}

static struct TGIDS *makeThreadGroup(int tgid)
{
  struct TGIDS *new_thread_group;
  /*
  printk(KERN_INFO "Make Thread Group for %d\n", tgid);
  */
  new_thread_group = vmalloc(sizeof(*new_thread_group));
  new_thread_group->tgid = tgid;

  INIT_LIST_HEAD(&new_thread_group->procs);
  
  return new_thread_group;
}

static struct ProcSeeds *makeProc(int pid)
{
  struct ProcSeeds *new_proc;

  printk(KERN_INFO "Make Proc for %d\n", pid);
  new_proc = vmalloc(sizeof(*new_proc));
  new_proc->pid = pid;
  INIT_LIST_HEAD(&new_proc->list);
  
  return new_proc;
}

static void addProc(struct TGIDS *thread_group, struct ProcSeeds *proc)
{
  struct ProcSeeds *prev;

  /*printk(KERN_INFO "Adding New Proc! PID: %d\n", proc->pid);
   */
  /* add this proc */
  list_add_tail(&proc->list, &thread_group->procs);

  /* go to previous, reseed from there */
  prev = list_entry(proc->list.prev, struct ProcSeeds, list);

  ReseedProcList(thread_group, prev->seed);
}

static struct TGIDS *getTGID(int tgid)
{
  struct list_head *node;
  struct TGIDS *thread_group;

  list_for_each(node, &thread_groups)
    {
      thread_group = list_entry(node, struct TGIDS, list);
      if(thread_group->tgid == tgid)
	{
	  /*
  printk(KERN_INFO "Thread group %d found\n", tgid);
	  */
	  return thread_group;
	}
    }
  /*
  printk(KERN_INFO "Thread group %d not found\n", tgid);
  */
  return NULL;
}

static struct ProcSeeds *getProc(struct TGIDS *thread_group, int pid)
{
  struct list_head *node;
  struct ProcSeeds *proc;

  list_for_each(node, &thread_group->procs)
    {
      proc = list_entry(node, struct ProcSeeds, list);
      if(proc->pid == pid)
	{
	  /*
	  printk(KERN_INFO "Proc %d found\n", pid);
	  */
	  return proc;
	}
    }
  /*
  printk(KERN_INFO "Proc %d not found\n", pid);
  */
  return NULL;
}

static void dumpProcList(struct TGIDS *thread_group)
{
  struct list_head *node;
  struct ProcSeeds *proc;

  printk(KERN_INFO "PROCLIST FOR THREAD GROUP %d\n", thread_group->tgid);
  list_for_each(node, &thread_group->procs)
    {
      proc = list_entry(node, struct ProcSeeds, list);
      printk(KERN_INFO "Proc: %d Seed: %lld\n", proc->pid, proc->seed);
    }
  printk(KERN_INFO "END DUMP\n");
}

static void ReseedProcList(struct TGIDS *thread_group, long long seed)
{
  struct list_head *node;
  struct ProcSeeds *curr_proc;

  
  printk(KERN_INFO "RESEEDING PROCLIST\n");
  printk(KERN_INFO "Pregame Dump:\n");
  
  dumpProcList(thread_group);
  

  list_for_each(node, &thread_group->procs)
    {
      curr_proc = list_entry(node, struct ProcSeeds, list);
      curr_proc->seed = seed;
      printk(KERN_INFO "Proc %d Assigned Seed %lld\n", curr_proc->pid, curr_proc->seed);
      
      seed = getRand(seed);
    }

  
  printk(KERN_INFO "Postgame Dump:\n");
  dumpProcList(thread_group);
  
}

static long long getRand(long long seed)
{
  return (MULTIPLIER*seed)%PMOD;
  /* Great for Debugging! */
  /*
    return (seed+1);
  */
}

static long long getProcsRandom(struct ProcSeeds *proc)
{
  struct list_head *node;
  struct ProcSeeds *curr_proc;
  struct ProcSeeds *next;
  long long nextSeed;

  /*
  printk(KERN_INFO "GETTING PROCESS RANDOM\n");
  */
  nextSeed = proc->seed;

  /*
  printk(KERN_INFO "Starting Process: %d, Starting Process' Seed: %lld\n", proc->pid, proc->seed);
  */

  next = list_entry(proc->list.next, struct ProcSeeds, list);

  /*
  printk(KERN_INFO "PROCLIST DUMP:\n");
  */

  list_for_each(node, &proc->list)
    {
      curr_proc = list_entry(node, struct ProcSeeds, list);
      /*
printk(KERN_INFO "Proc:%d has seed:%lld\n", curr_proc->pid, curr_proc->seed);
      */

    }
  /*
  printk(KERN_INFO "DONE\n");
  printk(KERN_INFO "Calculating Random:\n");
  */

  list_for_each(node, &next->list)
    {
      curr_proc = list_entry(node, struct ProcSeeds, list);
      nextSeed = getRand(nextSeed);
      /*  printk(KERN_INFO "At Proc %d, updated to %lld\n", curr_proc->pid, nextSeed);
       */
    }

  /*
  printk(KERN_INFO "ASSIGNING PROC %d SEED %lld\n", proc->pid, nextSeed);
  */
  proc->seed = nextSeed;
 

  /* printk(KERN_INFO "Returning Random Number %lld\n", nextSeed);
   */
  return nextSeed;
}


/* Module Macros */
module_init(lfprng_init);
module_exit(lfprng_exit);
